{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/fashion_mnist.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "N6ZDpd9XzFeN"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "KUu4vOt5zI9d"
      },
      "outputs": [],
      "source": [
        "# Copyright 2018 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "edfbxDDh2AEs"
      },
      "source": [
        "## Fashion MNIST with Keras and TPUs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RNo1Vfghpa8j"
      },
      "source": [
        "## Overview\n",
        "\n",
        "In this example, you can try out using tf.keras and Cloud TPUs to train a model on the fashion MNIST dataset. The model trains for 10 epochs on Cloud TPU and takes approximately 2 minutes to run.\n",
        "\n",
        "This notebook is hosted on GitHub. To view it in its original repository, after opening the notebook, select **File \u003e View on GitHub**."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dgAHfQtuhddd"
      },
      "source": [
        "## Learning objectives\n",
        "\n",
        "In this Colab, you will learn how to:\n",
        "*   Code for a standard conv-net that has 3 layers with drop-out and batch normalization between each layer in Keras.\n",
        "*   Create and compile the model under a distribution strategy in order ot use TPUs.\n",
        "*   Run a prediction to see how well the model can predict fashion categories and output the result."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "QrprJD-R-410"
      },
      "source": [
        "## Instructions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_I0RdnOSkNmi"
      },
      "source": [
        "\u003ch3\u003e  \u0026nbsp;\u0026nbsp;Train on TPU\u0026nbsp;\u0026nbsp; \u003ca href=\"https://cloud.google.com/tpu/\"\u003e\u003cimg valign=\"middle\" src=\"https://raw.githubusercontent.com/GoogleCloudPlatform/tensorflow-without-a-phd/master/tensorflow-rl-pong/images/tpu-hexagon.png\" width=\"50\"\u003e\u003c/a\u003e\u003c/h3\u003e\n",
        "\n",
        "1. On the main menu, click Runtime and select **Change runtime type**. Set \"TPU\" as the hardware accelerator.\n",
        "1. Click Runtime again and select **Runtime \u003e Run All**. You can also run the cells manually with Shift-ENTER. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5eEM-XOvURoU"
      },
      "source": [
        "TPUs are located in Google Cloud, for optimal performance, they read data directly from Google Cloud Storage (GCS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Lvo0t7XVIkWZ"
      },
      "source": [
        "## Data, model, and training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MICrRv8rmXVq"
      },
      "source": [
        "Begin by downloading the fashion MNIST dataset using `tf.keras.datasets`, as shown below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Zo-Yk6LFGfSf"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "\n",
        "import distutils\n",
        "if distutils.version.LooseVersion(tf.__version__) < '1.14':\n",
        "    raise Exception('This notebook is compatible with TensorFlow 1.14 or higher, for TensorFlow 1.13 or lower please use the previous version at https://github.com/tensorflow/tpu/blob/r1.13/tools/colab/fashion_mnist.ipynb')\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()\n",
        "\n",
        "# add empty color dimension\n",
        "x_train = np.expand_dims(x_train, -1)\n",
        "x_test = np.expand_dims(x_test, -1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Hgc2FZKVMx15"
      },
      "source": [
        "### Define the model\n",
        "\n",
        "The following example uses a standard conv-net that has 3 layers with drop-out and batch normalization between each layer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "W7gMbs70GxA7"
      },
      "outputs": [],
      "source": [
        "def create_model():\n",
        "  model = tf.keras.models.Sequential()\n",
        "  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))\n",
        "  model.add(tf.keras.layers.Conv2D(64, (5, 5), padding='same', activation='elu'))\n",
        "  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))\n",
        "  model.add(tf.keras.layers.Dropout(0.25))\n",
        "\n",
        "  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))\n",
        "  model.add(tf.keras.layers.Conv2D(128, (5, 5), padding='same', activation='elu'))\n",
        "  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))\n",
        "  model.add(tf.keras.layers.Dropout(0.25))\n",
        "\n",
        "  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))\n",
        "  model.add(tf.keras.layers.Conv2D(256, (5, 5), padding='same', activation='elu'))\n",
        "  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))\n",
        "  model.add(tf.keras.layers.Dropout(0.25))\n",
        "\n",
        "  model.add(tf.keras.layers.Flatten())\n",
        "  model.add(tf.keras.layers.Dense(256))\n",
        "  model.add(tf.keras.layers.Activation('elu'))\n",
        "  model.add(tf.keras.layers.Dropout(0.5))\n",
        "  model.add(tf.keras.layers.Dense(10))\n",
        "  model.add(tf.keras.layers.Activation('softmax'))\n",
        "  return model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xLeZATVaNAnE"
      },
      "source": [
        "### Train on the TPU\n",
        "\n",
        "To begin training, construct the model on the TPU and then compile it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "pWEYmd_hIWg8"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "resolver = tf.contrib.cluster_resolver.TPUClusterResolver('grpc://' + os.environ['COLAB_TPU_ADDR'])\n",
        "tf.contrib.distribute.initialize_tpu_system(resolver)\n",
        "strategy = tf.contrib.distribute.TPUStrategy(resolver)\n",
        "\n",
        "with strategy.scope():\n",
        "  model = create_model()\n",
        "  model.compile(\n",
        "      optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3, ),\n",
        "      loss='sparse_categorical_crossentropy',\n",
        "      metrics=['sparse_categorical_accuracy'])\n",
        "\n",
        "model.fit(\n",
        "    x_train.astype(np.float32), y_train.astype(np.float32),\n",
        "    epochs=17,\n",
        "    steps_per_epoch=60,\n",
        "    validation_data=(x_test.astype(np.float32), y_test.astype(np.float32)),\n",
        "    validation_freq=17\n",
        ")\n",
        "\n",
        "model.save_weights('./fashion_mnist.h5', overwrite=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ESL6ltQTMm05"
      },
      "source": [
        "### Check the results (inference)\n",
        "\n",
        "Now that you are done training, see how well the model can predict fashion categories!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "SaYPv_aKId2d"
      },
      "outputs": [],
      "source": [
        "LABEL_NAMES = ['t_shirt', 'trouser', 'pullover', 'dress', 'coat', 'sandal', 'shirt', 'sneaker', 'bag', 'ankle_boots']\n",
        "\n",
        "\n",
        "cpu_model = create_model()\n",
        "cpu_model.load_weights('./fashion_mnist.h5')\n",
        "\n",
        "from matplotlib import pyplot\n",
        "%matplotlib inline\n",
        "\n",
        "def plot_predictions(images, predictions):\n",
        "  n = images.shape[0]\n",
        "  nc = int(np.ceil(n / 4))\n",
        "  f, axes = pyplot.subplots(nc, 4)\n",
        "  for i in range(nc * 4):\n",
        "    y = i // 4\n",
        "    x = i % 4\n",
        "    axes[x, y].axis('off')\n",
        "    \n",
        "    label = LABEL_NAMES[np.argmax(predictions[i])]\n",
        "    confidence = np.max(predictions[i])\n",
        "    if i \u003e n:\n",
        "      continue\n",
        "    axes[x, y].imshow(images[i])\n",
        "    axes[x, y].text(0.5, 0.5, label + '\\n%.3f' % confidence, fontsize=14)\n",
        "\n",
        "  pyplot.gcf().set_size_inches(8, 8)  \n",
        "\n",
        "plot_predictions(np.squeeze(x_test[:16]), \n",
        "                 cpu_model.predict(x_test[:16]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "2a5cGsSTEBQD"
      },
      "source": [
        "## What's next\n",
        "\n",
        "* Learn about [Cloud TPUs](https://cloud.google.com/tpu/docs) that Google designed and optimized specifically to speed up and scale up ML workloads for training and inference and to enable ML engineers and researchers to iterate more quickly.\n",
        "* Explore the range of [Cloud TPU tutorials and Colabs](https://cloud.google.com/tpu/docs/tutorials) to find other examples that can be used when implementing your ML project.\n",
        "\n",
        "On Google Cloud Platform, in addition to GPUs and TPUs available on pre-configured [deep learning VMs](https://cloud.google.com/deep-learning-vm/),  you will find [AutoML](https://cloud.google.com/automl/)*(beta)* for training custom models without writing code and [Cloud ML Engine](https://cloud.google.com/ml-engine/docs/) which will allows you to run parallel trainings and hyperparameter tuning of your custom models on powerful distributed hardware.\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [
        "N6ZDpd9XzFeN"
      ],
      "name": "Keras Fashion MNIST",
      "provenance": [],
      "toc_visible": true,
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
